from typing import Optional
import datetime
import typer
from pathlib import Path
from functools import wraps
from rich.console import Console
from rich.panel import Panel
from rich.spinner import Spinner
from rich.live import Live
from rich.columns import Columns
from rich.markdown import Markdown
from rich.layout import Layout
from rich.text import Text
from rich.live import Live
from rich.table import Table
from collections import deque
import time
from rich.tree import Tree
from rich import box
from rich.align import Align
from rich.rule import Rule

from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.default_config import DEFAULT_CONFIG
from cli.models import AnalystType
from cli.utils import *

console = Console()

app = typer.Typer(
    name="TradingAgents",
    help="TradingAgents CLI: 多智能体 LLM 金融交易框架",
    add_completion=True,  # 启用 shell 补全
)


# 创建一个有最大长度的 deque 来存储最近的消息
class MessageBuffer:
    def __init__(self, max_length=100):
        self.messages = deque(maxlen=max_length)
        self.tool_calls = deque(maxlen=max_length)
        self.current_report = None
        self.final_report = None  # 存储完整的最终报告
        self.agent_status = {
            # 分析师团队
            "市场分析师": "待定",
            "社交分析师": "待定",
            "新闻分析师": "待定",
            "基本面分析师": "待定",
            # 研究团队
            "看涨研究员": "待定",
            "看跌研究员": "待定",
            "研究经理": "待定",
            # 交易团队
            "交易员": "待定",
            # 风险管理团队
            "风险分析师": "待定",
            "中立分析师": "待定",
            "保守分析师": "待定",
            # 投资组合管理团队
            "投资组合经理": "待定",
        }
        self.current_agent = None
        self.report_sections = {
            "market_report": None,
            "sentiment_report": None,
            "news_report": None,
            "fundamentals_report": None,
            "investment_plan": None,
            "trader_investment_plan": None,
            "final_trade_decision": None,
        }

    def add_message(self, message_type, content):
        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
        self.messages.append((timestamp, message_type, content))

    def add_tool_call(self, tool_name, args):
        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
        self.tool_calls.append((timestamp, tool_name, args))

    def update_agent_status(self, agent, status):
        if agent in self.agent_status:
            self.agent_status[agent] = status
            self.current_agent = agent

    def update_report_section(self, section_name, content):
        if section_name in self.report_sections:
            self.report_sections[section_name] = content
            self._update_current_report()

    def _update_current_report(self):
        # 对于面板显示，只显示最近更新的部分
        latest_section = None
        latest_content = None

        # 找到最近更新的部分
        for section, content in self.report_sections.items():
            if content is not None:
                latest_section = section
                latest_content = content
               
        if latest_section and latest_content:
            # 格式化当前部分以供显示
            section_titles = {
                "market_report": "市场分析",
                "sentiment_report": "社交情绪",
                "news_report": "新闻分析",
                "fundamentals_report": "基本面分析",
                "investment_plan": "研究团队决策",
                "trader_investment_plan": "交易团队计划",
                "final_trade_decision": "投资组合管理决策",
            }
            self.current_report = (
                f"### {section_titles[latest_section]}\n{latest_content}"
            )

        # 更新完整的最终报告
        self._update_final_report()

    def _update_final_report(self):
        report_parts = []

        # 分析师团队报告
        if any(
            self.report_sections[section]
            for section in [
                "market_report",
                "sentiment_report",
                "news_report",
                "fundamentals_report",
            ]
        ):
            report_parts.append("## 分析师团队报告")
            if self.report_sections["market_report"]:
                report_parts.append(
                    f"### 市场分析\n{self.report_sections['market_report']}"
                )
            if self.report_sections["sentiment_report"]:
                report_parts.append(
                    f"### 社交情绪\n{self.report_sections['sentiment_report']}"
                )
            if self.report_sections["news_report"]:
                report_parts.append(
                    f"### 新闻分析\n{self.report_sections['news_report']}"
                )
            if self.report_sections["fundamentals_report"]:
                report_parts.append(
                    f"### 基本面分析\n{self.report_sections['fundamentals_report']}"
                )

        # 研究团队报告
        if self.report_sections["investment_plan"]:
            report_parts.append("## 研究团队决策")
            report_parts.append(f"{self.report_sections['investment_plan']}")

        # 交易团队报告
        if self.report_sections["trader_investment_plan"]:
            report_parts.append("## 交易团队计划")
            report_parts.append(f"{self.report_sections['trader_investment_plan']}")

        # 投资组合管理决策
        if self.report_sections["final_trade_decision"]:
            report_parts.append("## 投资组合管理决策")
            report_parts.append(f"{self.report_sections['final_trade_decision']}")

        self.final_report = "\n\n".join(report_parts) if report_parts else None


message_buffer = MessageBuffer()


def create_layout():
    layout = Layout()
    layout.split_column(
        Layout(name="header", size=3),
        Layout(name="main"),
        Layout(name="footer", size=3),
    )
    layout["main"].split_column(
        Layout(name="upper", ratio=3), Layout(name="analysis", ratio=5)
    )
    layout["upper"].split_row(
        Layout(name="progress", ratio=2), Layout(name="messages", ratio=3)
    )
    return layout


def update_display(layout, spinner_text=None):
    # 带有欢迎信息的页眉
    layout["header"].update(
        Panel(
            "[bold green]欢迎使用 TradingAgents 命令行界面[/bold green]\n"
            "[dim]© [Tauric Research](https://github.com/TauricResearch)[/dim]",
            title="欢迎使用 TradingAgents",
            border_style="green",
            padding=(1, 2),
            expand=True,
        )
    )

    # 显示智能体状态的进度面板
    progress_table = Table(
        show_header=True,
        header_style="bold magenta",
        show_footer=False,
        box=box.SIMPLE_HEAD,  # 使用带水平线的简单页眉
        title=None,  # 移除多余的“进度”标题
        padding=(0, 2),  # 添加水平内边距
        expand=True,  # 使表格扩展以填充可用空间
    )
    progress_table.add_column("团队", style="cyan", justify="center", width=20)
    progress_table.add_column("智能体", style="green", justify="center", width=20)
    progress_table.add_column("状态", style="yellow", justify="center", width=20)

    # 按团队对智能体进行分组
    teams = {
        "分析师团队": [
            "市场分析师",
            "社交分析师",
            "新闻分析师",
            "基本面分析师",
        ],
        "研究团队": ["看涨研究员", "看跌研究员", "研究经理"],
        "交易团队": ["交易员"],
        "风险管理": ["风险分析师", "中立分析师", "保守分析师"],
        "投资组合管理": ["投资组合经理"],
    }

    for team, agents in teams.items():
        # 添加带有团队名称的第一个智能体
        first_agent = agents[0]
        status = message_buffer.agent_status[first_agent]
        if status == "进行中":
            spinner = Spinner(
                "dots", text="[blue]进行中[/blue]", style="bold cyan"
            )
            status_cell = spinner
        else:
            status_color = {
                "待定": "yellow",
                "已完成": "green",
                "错误": "red",
            }.get(status, "white")
            status_cell = f"[{status_color}]{status}[/{status_color}]"
        progress_table.add_row(team, first_agent, status_cell)

        # 添加团队中的其余智能体
        for agent in agents[1:]:
            status = message_buffer.agent_status[agent]
            if status == "进行中":
                spinner = Spinner(
                    "dots", text="[blue]进行中[/blue]", style="bold cyan"
                )
                status_cell = spinner
            else:
                status_color = {
                    "待定": "yellow",
                    "已完成": "green",
                    "错误": "red",
                }.get(status, "white")
                status_cell = f"[{status_color}]{status}[/{status_color}]"
            progress_table.add_row("", agent, status_cell)

        # 在每个团队后添加水平线
        progress_table.add_row("─" * 20, "─" * 20, "─" * 20, style="dim")

    layout["progress"].update(
        Panel(progress_table, title="进度", border_style="cyan", padding=(1, 2))
    )

    # 显示最近消息和工具调用的消息面板
    messages_table = Table(
        show_header=True,
        header_style="bold magenta",
        show_footer=False,
        expand=True,  # 使表格扩展以填充可用空间
        box=box.MINIMAL,  # 使用最小化的框样式以获得更轻量的外观
        show_lines=True,  # 保留水平线
        padding=(0, 1),  # 在列之间添加一些内边距
    )
    messages_table.add_column("时间", style="cyan", width=8, justify="center")
    messages_table.add_column("类型", style="green", width=10, justify="center")
    messages_table.add_column(
        "内容", style="white", no_wrap=False, ratio=1
    )  # 使内容列扩展

    # 合并工具调用和消息
    all_messages = []

    # 添加工具调用
    for timestamp, tool_name, args in message_buffer.tool_calls:
        # 如果工具调用参数太长，则截断
        if isinstance(args, str) and len(args) > 100:
            args = args[:97] + "..."
        all_messages.append((timestamp, "工具", f"{tool_name}: {args}"))

    # 添加常规消息
    for timestamp, msg_type, content in message_buffer.messages:
        # 如果内容不是字符串，则转换为字符串
        content_str = content
        if isinstance(content, list):
            # 处理内容块列表（Anthropic 格式）
            text_parts = []
            for item in content:
                if isinstance(item, dict):
                    if item.get('type') == 'text':
                        text_parts.append(item.get('text', ''))
                    elif item.get('type') == 'tool_use':
                        text_parts.append(f"[工具: {item.get('name', 'unknown')}]")
                else:
                    text_parts.append(str(item))
            content_str = ' '.join(text_parts)
        elif not isinstance(content_str, str):
            content_str = str(content)
            
        # 如果消息内容太长，则截断
        if len(content_str) > 200:
            content_str = content_str[:197] + "..."
        all_messages.append((timestamp, msg_type, content_str))

    # 按时间戳排序
    all_messages.sort(key=lambda x: x[0])

    # 根据可用空间计算可以显示多少条消息
    # 从一个合理的数字开始，并根据内容长度进行调整
    max_messages = 12  # 从 8 增加到 12 以更好地填充空间

    # 获取将适合面板的最后 N 条消息
    recent_messages = all_messages[-max_messages:]

    # 将消息添加到表格
    for timestamp, msg_type, content in recent_messages:
        # 格式化内容并自动换行
        wrapped_content = Text(content, overflow="fold")
        messages_table.add_row(timestamp, msg_type, wrapped_content)

    if spinner_text:
        messages_table.add_row("", "加载中", spinner_text)

    # 添加页脚以指示消息是否被截断
    if len(all_messages) > max_messages:
        messages_table.footer = (
            f"[dim]显示最近 {max_messages} 条消息，共 {len(all_messages)} 条[/dim]"
        )

    layout["messages"].update(
        Panel(
            messages_table,
            title="消息和工具",
            border_style="blue",
            padding=(1, 2),
        )
    )

    # 显示当前报告的分析面板
    if message_buffer.current_report:
        layout["analysis"].update(
            Panel(
                Markdown(message_buffer.current_report),
                title="当前报告",
                border_style="green",
                padding=(1, 2),
            )
        )
    else:
        layout["analysis"].update(
            Panel(
                "[italic]等待分析报告...[/italic]",
                title="当前报告",
                border_style="green",
                padding=(1, 2),
            )
        )

    # 带有统计信息的页脚
    tool_calls_count = len(message_buffer.tool_calls)
    llm_calls_count = sum(
        1 for _, msg_type, _ in message_buffer.messages if msg_type == "推理"
    )
    reports_count = sum(
        1 for content in message_buffer.report_sections.values() if content is not None
    )

    stats_table = Table(show_header=False, box=None, padding=(0, 2), expand=True)
    stats_table.add_column("统计", justify="center")
    stats_table.add_row(
        f"工具调用: {tool_calls_count} | LLM 调用: {llm_calls_count} | 生成的报告: {reports_count}"
    )

    layout["footer"].update(Panel(stats_table, border_style="grey50"))


def get_user_selections():
    """在开始分析显示之前获取所有用户选择。"""
    # 显示 ASCII 艺术欢迎信息
    with open("./cli/static/welcome.txt", "r") as f:
        welcome_ascii = f.read()

    # 创建欢迎框内容
    welcome_content = f"{welcome_ascii}\n"
    welcome_content += "[bold green]TradingAgents: 多智能体 LLM 金融交易框架 - 命令行界面[/bold green]\n\n"
    welcome_content += "[bold]工作流程步骤:[/bold]\n"
    welcome_content += "I. 分析师团队 → II. 研究团队 → III. 交易员 → IV. 风险管理 → V. 投资组合管理\n\n"
    welcome_content += (
        "[dim]由 [Tauric Research](https://github.com/TauricResearch) 构建[/dim]"
    )

    # 创建并居中欢迎框
    welcome_box = Panel(
        welcome_content,
        border_style="green",
        padding=(1, 2),
        title="欢迎使用 TradingAgents",
        subtitle="多智能体 LLM 金融交易框架",
    )
    console.print(Align.center(welcome_box))
    console.print()  # 在欢迎框后添加一个空行

    # 为每个步骤创建一个带框的问卷
    def create_question_box(title, prompt, default=None):
        box_content = f"[bold]{title}[/bold]\n"
        box_content += f"[dim]{prompt}[/dim]"
        if default:
            box_content += f"\n[dim]默认: {default}[/dim]"
        return Panel(box_content, border_style="blue", padding=(1, 2))

    # 步骤 1: 股票代码
    console.print(
        create_question_box(
            "步骤 1: 股票代码", "输入要分析的股票代码", "SPY"
        )
    )
    selected_ticker = get_ticker()

    # 步骤 2: 分析日期
    default_date = datetime.datetime.now().strftime("%Y-%m-%d")
    console.print(
        create_question_box(
            "步骤 2: 分析日期",
            "输入分析日期 (YYYY-MM-DD)",
            default_date,
        )
    )
    analysis_date = get_analysis_date()

    # 步骤 3: 选择分析师
    console.print(
        create_question_box(
            "步骤 3: 分析师团队", "为分析选择您的 LLM 分析师智能体"
        )
    )
    selected_analysts = select_analysts()
    console.print(
        f"[green]已选分析师:[/green] {', '.join(analyst.value for analyst in selected_analysts)}"
    )

    # 步骤 4: 研究深度
    console.print(
        create_question_box(
            "步骤 4: 研究深度", "选择您的研究深度级别"
        )
    )
    selected_research_depth = select_research_depth()

    # 步骤 5: OpenAI 后端
    console.print(
        create_question_box(
            "步骤 5: OpenAI 后端", "选择要与之通信的服务"
        )
    )
    selected_llm_provider, backend_url = select_llm_provider()
    
    # 步骤 6: 思维智能体
    console.print(
        create_question_box(
            "步骤 6: 思维智能体", "为分析选择您的思维智能体"
        )
    )
    selected_shallow_thinker = select_shallow_thinking_agent(selected_llm_provider)
    selected_deep_thinker = select_deep_thinking_agent(selected_llm_provider)

    return {
        "ticker": selected_ticker,
        "analysis_date": analysis_date,
        "analysts": selected_analysts,
        "research_depth": selected_research_depth,
        "llm_provider": selected_llm_provider.lower(),
        "backend_url": backend_url,
        "shallow_thinker": selected_shallow_thinker,
        "deep_thinker": selected_deep_thinker,
    }


def get_ticker():
    """从用户输入获取股票代码。"""
    return typer.prompt("", default="SPY")


def get_analysis_date():
    """从用户输入获取分析日期。"""
    while True:
        date_str = typer.prompt(
            "", default=datetime.datetime.now().strftime("%Y-%m-%d")
        )
        try:
            # 验证日期格式并确保不是未来日期
            analysis_date = datetime.datetime.strptime(date_str, "%Y-%m-%d")
            if analysis_date.date() > datetime.datetime.now().date():
                console.print("[red]错误: 分析日期不能是未来日期[/red]")
                continue
            return date_str
        except ValueError:
            console.print(
                "[red]错误: 日期格式无效。请使用 YYYY-MM-DD[/red]"
            )


def display_complete_report(final_state):
    """显示带有基于团队的面板的完整分析报告。"""
    console.print("\n[bold green]完整分析报告[/bold green]\n")

    # I. 分析师团队报告
    analyst_reports = []

    # 市场分析师报告
    if final_state.get("market_report"):
        analyst_reports.append(
            Panel(
                Markdown(final_state["market_report"]),
                title="市场分析师",
                border_style="blue",
                padding=(1, 2),
            )
        )

    # 社交分析师报告
    if final_state.get("sentiment_report"):
        analyst_reports.append(
            Panel(
                Markdown(final_state["sentiment_report"]),
                title="社交分析师",
                border_style="blue",
                padding=(1, 2),
            )
        )

    # 新闻分析师报告
    if final_state.get("news_report"):
        analyst_reports.append(
            Panel(
                Markdown(final_state["news_report"]),
                title="新闻分析师",
                border_style="blue",
                padding=(1, 2),
            )
        )

    # 基本面分析师报告
    if final_state.get("fundamentals_report"):
        analyst_reports.append(
            Panel(
                Markdown(final_state["fundamentals_report"]),
                title="基本面分析师",
                border_style="blue",
                padding=(1, 2),
            )
        )

    if analyst_reports:
        console.print(
            Panel(
                Columns(analyst_reports, equal=True, expand=True),
                title="I. 分析师团队报告",
                border_style="cyan",
                padding=(1, 2),
            )
        )

    # II. 研究团队报告
    if final_state.get("investment_debate_state"):
        research_reports = []
        debate_state = final_state["investment_debate_state"]

        # 看涨研究员分析
        if debate_state.get("bull_history"):
            research_reports.append(
                Panel(
                    Markdown(debate_state["bull_history"]),
                    title="看涨研究员",
                    border_style="blue",
                    padding=(1, 2),
                )
            )

        # 看跌研究员分析
        if debate_state.get("bear_history"):
            research_reports.append(
                Panel(
                    Markdown(debate_state["bear_history"]),
                    title="看跌研究员",
                    border_style="blue",
                    padding=(1, 2),
                )
            )

        # 研究经理决策
        if debate_state.get("judge_decision"):
            research_reports.append(
                Panel(
                    Markdown(debate_state["judge_decision"]),
                    title="研究经理",
                    border_style="blue",
                    padding=(1, 2),
                )
            )

        if research_reports:
            console.print(
                Panel(
                    Columns(research_reports, equal=True, expand=True),
                    title="II. 研究团队决策",
                    border_style="magenta",
                    padding=(1, 2),
                )
            )

    # III. 交易团队报告
    if final_state.get("trader_investment_plan"):
        console.print(
            Panel(
                Panel(
                    Markdown(final_state["trader_investment_plan"]),
                    title="交易员",
                    border_style="blue",
                    padding=(1, 2),
                ),
                title="III. 交易团队计划",
                border_style="yellow",
                padding=(1, 2),
            )
        )

    # IV. 风险管理团队报告
    if final_state.get("risk_debate_state"):
        risk_reports = []
        risk_state = final_state["risk_debate_state"]

        # 激进 (风险) 分析师分析
        if risk_state.get("risky_history"):
            risk_reports.append(
                Panel(
                    Markdown(risk_state["risky_history"]),
                    title="激进分析师",
                    border_style="blue",
                    padding=(1, 2),
                )
            )

        # 保守 (安全) 分析师分析
        if risk_state.get("safe_history"):
            risk_reports.append(
                Panel(
                    Markdown(risk_state["safe_history"]),
                    title="保守分析师",
                    border_style="blue",
                    padding=(1, 2),
                )
            )

        # 中立分析师分析
        if risk_state.get("neutral_history"):
            risk_reports.append(
                Panel(
                    Markdown(risk_state["neutral_history"]),
                    title="中立分析师",
                    border_style="blue",
                    padding=(1, 2),
                )
            )

        if risk_reports:
            console.print(
                Panel(
                    Columns(risk_reports, equal=True, expand=True),
                    title="IV. 风险管理团队决策",
                    border_style="red",
                    padding=(1, 2),
                )
            )

        # V. 投资组合经理决策
        if risk_state.get("judge_decision"):
            console.print(
                Panel(
                    Panel(
                        Markdown(risk_state["judge_decision"]),
                        title="投资组合经理",
                        border_style="blue",
                        padding=(1, 2),
                    ),
                    title="V. 投资组合经理决策",
                    border_style="green",
                    padding=(1, 2),
                )
            )


def update_research_team_status(status):
    """更新所有研究团队成员和交易员的状态。"""
    research_team = ["看涨研究员", "看跌研究员", "研究经理", "交易员"]
    for agent in research_team:
        message_buffer.update_agent_status(agent, status)

def extract_content_string(content):
    """从各种消息格式中提取字符串内容。"""
    if isinstance(content, str):
        return content
    elif isinstance(content, list):
        # 处理 Anthropic 的列表格式
        text_parts = []
        for item in content:
            if isinstance(item, dict):
                if item.get('type') == 'text':
                    text_parts.append(item.get('text', ''))
                elif item.get('type') == 'tool_use':
                    text_parts.append(f"[工具: {item.get('name', 'unknown')}]")
            else:
                text_parts.append(str(item))
        return ' '.join(text_parts)
    else:
        return str(content)

def run_analysis():
    # 首先获取所有用户选择
    selections = get_user_selections()

    # 使用选定的研究深度创建配置
    config = DEFAULT_CONFIG.copy()
    config["max_debate_rounds"] = selections["research_depth"]
    config["max_risk_discuss_rounds"] = selections["research_depth"]
    config["quick_think_llm"] = selections["shallow_thinker"]
    config["deep_think_llm"] = selections["deep_thinker"]
    config["backend_url"] = selections["backend_url"]
    config["llm_provider"] = selections["llm_provider"].lower()

    # 初始化图
    graph = TradingAgentsGraph(
        [analyst.value for analyst in selections["analysts"]], config=config, debug=True
    )

    # 创建结果目录
    results_dir = Path(config["results_dir"]) / selections["ticker"] / selections["analysis_date"]
    results_dir.mkdir(parents=True, exist_ok=True)
    report_dir = results_dir / "reports"
    report_dir.mkdir(parents=True, exist_ok=True)
    log_file = results_dir / "message_tool.log"
    log_file.touch(exist_ok=True)

    def save_message_decorator(obj, func_name):
        func = getattr(obj, func_name)
        @wraps(func)
        def wrapper(*args, **kwargs):
            func(*args, **kwargs)
            timestamp, message_type, content = obj.messages[-1]
            content = content.replace("\n", " ")  # 将换行符替换为空格
            with open(log_file, "a") as f:
                f.write(f"{timestamp} [{message_type}] {content}\n")
        return wrapper
    
    def save_tool_call_decorator(obj, func_name):
        func = getattr(obj, func_name)
        @wraps(func)
        def wrapper(*args, **kwargs):
            func(*args, **kwargs)
            timestamp, tool_name, args = obj.tool_calls[-1]
            args_str = ", ".join(f"{k}={v}" for k, v in args.items())
            with open(log_file, "a") as f:
                f.write(f"{timestamp} [工具调用] {tool_name}({args_str})\n")
        return wrapper

    def save_report_section_decorator(obj, func_name):
        func = getattr(obj, func_name)
        @wraps(func)
        def wrapper(section_name, content):
            func(section_name, content)
            if section_name in obj.report_sections and obj.report_sections[section_name] is not None:
                content = obj.report_sections[section_name]
                if content:
                    file_name = f"{section_name}.md"
                    with open(report_dir / file_name, "w") as f:
                        f.write(content)
        return wrapper

    message_buffer.add_message = save_message_decorator(message_buffer, "add_message")
    message_buffer.add_tool_call = save_tool_call_decorator(message_buffer, "add_tool_call")
    message_buffer.update_report_section = save_report_section_decorator(message_buffer, "update_report_section")

    # 现在开始显示布局
    layout = create_layout()

    with Live(layout, refresh_per_second=4) as live:
        # 初始显示
        update_display(layout)

        # 添加初始消息
        message_buffer.add_message("系统", f"选定股票代码: {selections['ticker']}")
        message_buffer.add_message(
            "系统", f"分析日期: {selections['analysis_date']}"
        )
        message_buffer.add_message(
            "系统",
            f"选定分析师: {', '.join(analyst.value for analyst in selections['analysts'])}",
        )
        update_display(layout)

        # 重置智能体状态
        for agent in message_buffer.agent_status:
            message_buffer.update_agent_status(agent, "待定")

        # 重置报告部分
        for section in message_buffer.report_sections:
            message_buffer.report_sections[section] = None
        message_buffer.current_report = None
        message_buffer.final_report = None

        # 将第一个分析师的状态更新为进行中
        first_analyst = f"{selections['analysts'][0].value.capitalize()} 分析师"
        message_buffer.update_agent_status(first_analyst, "进行中")
        update_display(layout)

        # 创建加载文本
        spinner_text = (
            f"正在分析 {selections['ticker']} 在 {selections['analysis_date']}..."
        )
        update_display(layout, spinner_text)

        # 初始化状态并获取图参数
        init_agent_state = graph.propagator.create_initial_state(
            selections["ticker"], selections["analysis_date"]
        )
        args = graph.propagator.get_graph_args()

        # 流式传输分析
        trace = []
        for chunk in graph.graph.stream(init_agent_state, **args):
            if len(chunk["messages"]) > 0:
                # 获取块中的最后一条消息
                last_message = chunk["messages"][-1]

                # 提取消息内容和类型
                if hasattr(last_message, "content"):
                    content = extract_content_string(last_message.content)  # 使用辅助函数
                    msg_type = "推理"
                else:
                    content = str(last_message)
                    msg_type = "系统"

                # 将消息添加到缓冲区
                message_buffer.add_message(msg_type, content)                

                # 如果是工具调用，则添加到工具调用中
                if hasattr(last_message, "tool_calls"):
                    for tool_call in last_message.tool_calls:
                        # 处理字典和对象工具调用
                        if isinstance(tool_call, dict):
                            message_buffer.add_tool_call(
                                tool_call["name"], tool_call["args"]
                            )
                        else:
                            message_buffer.add_tool_call(tool_call.name, tool_call.args)

                # 根据块内容更新报告和智能体状态
                # 分析师团队报告
                if "market_report" in chunk and chunk["market_report"]:
                    message_buffer.update_report_section(
                        "market_report", chunk["market_report"]
                    )
                    message_buffer.update_agent_status("市场分析师", "已完成")
                    # 将下一个分析师设置为进行中
                    if "social" in selections["analysts"]:
                        message_buffer.update_agent_status(
                            "社交分析师", "进行中"
                        )

                if "sentiment_report" in chunk and chunk["sentiment_report"]:
                    message_buffer.update_report_section(
                        "sentiment_report", chunk["sentiment_report"]
                    )
                    message_buffer.update_agent_status("社交分析师", "已完成")
                    # 将下一个分析师设置为进行中
                    if "news" in selections["analysts"]:
                        message_buffer.update_agent_status(
                            "新闻分析师", "进行中"
                        )

                if "news_report" in chunk and chunk["news_report"]:
                    message_buffer.update_report_section(
                        "news_report", chunk["news_report"]
                    )
                    message_buffer.update_agent_status("新闻分析师", "已完成")
                    # 将下一个分析师设置为进行中
                    if "fundamentals" in selections["analysts"]:
                        message_buffer.update_agent_status(
                            "基本面分析师", "进行中"
                        )

                if "fundamentals_report" in chunk and chunk["fundamentals_report"]:
                    message_buffer.update_report_section(
                        "fundamentals_report", chunk["fundamentals_report"]
                    )
                    message_buffer.update_agent_status(
                        "基本面分析师", "已完成"
                    )
                    # 将所有研究团队成员设置为进行中
                    update_research_team_status("进行中")

                # 研究团队 - 处理投资辩论状态
                if (
                    "investment_debate_state" in chunk
                    and chunk["investment_debate_state"]
                ):
                    debate_state = chunk["investment_debate_state"]

                    # 更新看涨研究员状态和报告
                    if "bull_history" in debate_state and debate_state["bull_history"]:
                        # 保持所有研究团队成员进行中
                        update_research_team_status("进行中")
                        # 提取最新的看涨回应
                        bull_responses = debate_state["bull_history"].split("\n")
                        latest_bull = bull_responses[-1] if bull_responses else ""
                        if latest_bull:
                            message_buffer.add_message("推理", latest_bull)
                            # 使用看涨研究员的最新分析更新研究报告
                            message_buffer.update_report_section(
                                "investment_plan",
                                f"### 看涨研究员分析\n{latest_bull}",
                            )

                    # 更新看跌研究员状态和报告
                    if "bear_history" in debate_state and debate_state["bear_history"]:
                        # 保持所有研究团队成员进行中
                        update_research_team_status("进行中")
                        # 提取最新的看跌回应
                        bear_responses = debate_state["bear_history"].split("\n")
                        latest_bear = bear_responses[-1] if bear_responses else ""
                        if latest_bear:
                            message_buffer.add_message("推理", latest_bear)
                            # 使用看跌研究员的最新分析更新研究报告
                            message_buffer.update_report_section(
                                "investment_plan",
                                f"{message_buffer.report_sections['investment_plan']}\n\n### 看跌研究员分析\n{latest_bear}",
                            )

                    # 更新研究经理状态和最终决策
                    if (
                        "judge_decision" in debate_state
                        and debate_state["judge_decision"]
                    ):
                        # 在最终决策之前保持所有研究团队成员进行中
                        update_research_team_status("进行中")
                        message_buffer.add_message(
                            "推理",
                            f"研究经理: {debate_state['judge_decision']}",
                        )
                        # 使用最终决策更新研究报告
                        message_buffer.update_report_section(
                            "investment_plan",
                            f"{message_buffer.report_sections['investment_plan']}\n\n### 研究经理决策\n{debate_state['judge_decision']}",
                        )
                        # 将所有研究团队成员标记为已完成
                        update_research_team_status("已完成")
                        # 将第一个风险分析师设置为进行中
                        message_buffer.update_agent_status(
                            "风险分析师", "进行中"
                        )

                # 交易团队
                if (
                    "trader_investment_plan" in chunk
                    and chunk["trader_investment_plan"]
                ):
                    message_buffer.update_report_section(
                        "trader_investment_plan", chunk["trader_investment_plan"]
                    )
                    # 将第一个风险分析师设置为进行中
                    message_buffer.update_agent_status("风险分析师", "进行中")

                # 风险管理团队 - 处理风险辩论状态
                if "risk_debate_state" in chunk and chunk["risk_debate_state"]:
                    risk_state = chunk["risk_debate_state"]

                    # 更新风险分析师状态和报告
                    if (
                        "current_risky_response" in risk_state
                        and risk_state["current_risky_response"]
                    ):
                        message_buffer.update_agent_status(
                            "风险分析师", "进行中"
                        )
                        message_buffer.add_message(
                            "推理",
                            f"风险分析师: {risk_state['current_risky_response']}",
                        )
                        # 仅使用风险分析师的最新分析更新风险报告
                        message_buffer.update_report_section(
                            "final_trade_decision",
                            f"### 风险分析师分析\n{risk_state['current_risky_response']}",
                        )

                    # 更新保守分析师状态和报告
                    if (
                        "current_safe_response" in risk_state
                        and risk_state["current_safe_response"]
                    ):
                        message_buffer.update_agent_status(
                            "保守分析师", "进行中"
                        )
                        message_buffer.add_message(
                            "推理",
                            f"保守分析师: {risk_state['current_safe_response']}",
                        )
                        # 仅使用保守分析师的最新分析更新风险报告
                        message_buffer.update_report_section(
                            "final_trade_decision",
                            f"### 保守分析师分析\n{risk_state['current_safe_response']}",
                        )

                    # 更新中立分析师状态和报告
                    if (
                        "current_neutral_response" in risk_state
                        and risk_state["current_neutral_response"]
                    ):
                        message_buffer.update_agent_status(
                            "中立分析师", "进行中"
                        )
                        message_buffer.add_message(
                            "推理",
                            f"中立分析师: {risk_state['current_neutral_response']}",
                        )
                        # 仅使用中立分析师的最新分析更新风险报告
                        message_buffer.update_report_section(
                            "final_trade_decision",
                            f"### 中立分析师分析\n{risk_state['current_neutral_response']}",
                        )

                    # 更新投资组合经理状态和最终决策
                    if "judge_decision" in risk_state and risk_state["judge_decision"]:
                        message_buffer.update_agent_status(
                            "投资组合经理", "进行中"
                        )
                        message_buffer.add_message(
                            "推理",
                            f"投资组合经理: {risk_state['judge_decision']}",
                        )
                        # 仅使用最终决策更新风险报告
                        message_buffer.update_report_section(
                            "final_trade_decision",
                            f"### 投资组合经理决策\n{risk_state['judge_decision']}",
                        )
                        # 将风险分析师标记为已完成
                        message_buffer.update_agent_status("风险分析师", "已完成")
                        message_buffer.update_agent_status("保守分析师", "已完成")
                        message_buffer.update_agent_status(
                            "中立分析师", "已完成"
                        )
                        message_buffer.update_agent_status(
                            "投资组合经理", "已完成"
                        )

                # 更新显示
                update_display(layout)

            trace.append(chunk)

        # 获取最终状态和决策
        final_state = trace[-1]
        decision = graph.process_signal(final_state["final_trade_decision"])

        # 将所有智能体状态更新为已完成
        for agent in message_buffer.agent_status:
            message_buffer.update_agent_status(agent, "已完成")

        message_buffer.add_message(
            "分析", f"已完成对 {selections['analysis_date']} 的分析"
        )

        # 更新最终报告部分
        for section in message_buffer.report_sections.keys():
            if section in final_state:
                message_buffer.update_report_section(section, final_state[section])

        # 显示完整的最终报告
        display_complete_report(final_state)

        update_display(layout)


@app.command()
def analyze():
    run_analysis()


if __name__ == "__main__":
    app()